#include <application.hpp>

const std::map<int ,std::string> robot_fsm_map = {
    {1, "Damp"},
    {4, "StandUp"},
    {2, "Squat"},
    {3, "Sit"},
    {0, "ZeroTorque"},
    {500, "Start"},
};

const std::string QUEUE_NAME("/robotctl_msg");
const std::string SERVER_ADDRESS("139.155.134.111:1883");
const std::string USER_NAME("yufeng");
const std::string PASSWORD("xuanye123");
const std::string CLIENT_ID("robot_123");                   // 客户端 ID
const std::string TOPIC("robot/#");                        // 主题
const int QOS = 1;                                          // 消息质量等级

Application::Application(int argc, char *argv[]) {
    Initialize();
}

Application::~Application() {
    mqttclient_->disconnect()->wait();
    radar_->Stop();
}

void Application::Initialize() {
    radar_ = radar::GetInstance();
    radar_->Init();
    radar_->KalmanInit();
    radar_->setFollowDistance(100.0);
    radar_->setFollowAngle(0.0);
    radar_->setAngleDeviation(5);
    radar_->setDistanceDeviation(10);

    unitree::robot::ChannelFactory::Instance()->Init(0, "eth0");
    std::cout << "ChannelFactory Init!" << std::endl;
    roclient_ = std::make_unique<unitree::robot::g1::LocoClient>();
    armclient_ = std::make_unique<unitree::robot::g1::G1ArmActionClient>();
    mqttclient_ = std::make_unique<mqtt::async_client>(SERVER_ADDRESS, CLIENT_ID);
    mqttcb_ = std::make_unique<callback>(*mqttclient_, options_, TOPIC, QOS);

    roclient_->Init();
    roclient_->SetTimeout(10.f);

    armclient_->Init();
    armclient_->SetTimeout(10.f);

    int mode;
    roclient_->GetFsmMode(mode);
    std::cout << "Robot mode: " << mode << std::endl;
    
    audio_.Init();
    coordinates_.Init();

    audio_.setVolume(80);
    
    options_.set_keep_alive_interval(20);
    options_.set_clean_session(true);
    options_.set_user_name(USER_NAME);
    options_.set_password(PASSWORD);

    mqttclient_->set_callback(*mqttcb_);

    mq_ = mq_open(QUEUE_NAME.c_str(), O_RDONLY | O_NONBLOCK);
    if (mq_ == (mqd_t)-1) {
        std::cerr << "Failed to open message queue" << std::endl;
        return;
    }

    if (mq_getattr(mq_, &attr_) == -1) {
        std::cerr << "Failed to get message queue attributes" << std::endl;
        mq_close(mq_);
        return;
    }

    msg_buffer.resize(attr_.mq_msgsize);

    
}

void Application::Loop() {
    while (true) {
        std::list<std::function<void()>> task_;
        {
            std::unique_lock<std::mutex> lock(mutex_);
            cv_.wait(lock);
            task_ = std::move(main_task_);
        }
        for(auto& task : task_) {
            task();
        }
    }
}

void Application::Schedule(std::function<void()> func) {
    {
        std::unique_lock<std::mutex> lock(mutex_);
        cv_.notify_one();
        main_task_.push_back(std::move(func));
    }
}

void Application::POSIX_Message() {
    if (mq_ != -1) {
        while (true) {
            try {
                ssize_t bytes_read = mq_receive(mq_, msg_buffer.data(), attr_.mq_msgsize, nullptr);
                if (bytes_read >= 0) {
                    std::string msg(msg_buffer.data(), bytes_read);
                    std::cout << "Received message: " << msg << std::endl;
                    auto robotctl = nlohmann::json::parse(msg);
                    if (robotctl.contains("wake_up")) {
                        if (robotctl["wake_up"] == "true") {
                            {
                                std::unique_lock<std::mutex> lock(audio_.led_mutex_);
                                audio_.wakeup_.store(true);
                                audio_.led_cv_.notify_one();
                            }
                        }
                        else if (robotctl["wake_up"] == "false") {
                            {
                                std::unique_lock<std::mutex> lock(audio_.led_mutex_);
                                audio_.wakeup_.store(false);
                                audio_.led_cv_.notify_one();
                            }
                        }
                    }
                    else if (robotctl.contains("audio_merge")) {
                        if(robotctl["audio_merge"] == "Dialogue") {
                            {
                                std::unique_lock<std::mutex> lock(audio_.audio_mutex_);
                                audio_.playsta.store(Play_Dialogue);
                                audio_.audio_cv_.notify_one();
                            }
                        }
                        else if(robotctl["audio_merge"] == "Music") {
                            {
                                std::unique_lock<std::mutex> lock(audio_.audio_mutex_);
                                audio_.playsta.store(Play_Music);
                                audio_.audio_cv_.notify_one();
                            }
                        } 
                        else if(robotctl["audio_merge"] == "Stop") {
                            {
                                std::unique_lock<std::mutex> lock(audio_.audio_mutex_);
                                audio_.playsta.store(Play_Stop);
                                audio_.audio_cv_.notify_one();
                            }
                        }
                    }

                    else if (robotctl.contains("exit")) {
                        roclient_->Damp();
                        break;
                    }
                } 
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
            }
            catch (const nlohmann::json::parse_error& e) {
                std::cerr << "JSON parse error: " << e.what() << std::endl;
            }
            catch (const std::exception& exc) {
                std::cerr << "Error: " << exc.what() << std::endl;
            }
        }
    }
}

void Application::MQTT_Message() {
    while(mqttcb_->Is_Running()) {
        try {
            // 读取 Unitree 机器人的数据
            std::pair<std::string, std::string> msg;
            {
                std::unique_lock<std::mutex> lock(mqttcb_->queue_mutex);
                mqttcb_->queue_cv.wait(lock, [&] { return !mqttcb_->msg_queue.empty() || !mqttcb_->Is_Running(); });
                if (!mqttcb_->Is_Running()) break;
                if (!mqttcb_->msg_queue.empty()) {
                    msg = mqttcb_->msg_queue.front();
                    mqttcb_->msg_queue.pop();
                }   
            }
            
            Parse_Json(msg.second);
        }
        catch (const nlohmann::json::parse_error& e) {
            std::cerr << "JSON parse error: " << e.what() << std::endl;
        }
        catch (const std::exception& exc) {
            std::cerr << "Error: " << exc.what() << std::endl;
        }
    }
}

int Application::run() {
    while (true) {
        try {
            mqttclient_->connect(options_, nullptr, *mqttcb_)->wait();
            std::cout << "Connected to MQTT broker." << std::endl;

            radar_->OnReceiveData([this](std::vector<uint8_t>& buffer) {
                int head = buffer[0] | (buffer[1] << 8) | (buffer[2] << 16) | (buffer[3] << 24);
                
                if (0xFFFFFFFF != head) {
                    std::cout << "Radar data valid error!" << std::endl;
                    return;
                }

                int cmd_id = (buffer[8] << 8) | buffer[9];
                if (cmd_id == 0x2001) {
                    this->radar_->buffer_to_radar_frame(buffer);
                    if (!this->radar_->GetDetected())
                        this->radar_->SetDetected(true);
                }
                else {
                    if (this->radar_->GetDetected())
                        this->radar_->SetDetected(false);
                    // this->radar_->follow_mode_.store(false);
                    // std::cout << "non detected!" << std::endl;
                }

                if (this->radar_->follow_mode_.load() && this->radar_->GetDetected()) {
                    auto radar_frame = this->radar_->GetRadarFrame();
                    float x_vel{0};
                    float z_vel{0};
                    if (std::abs(this->radar_->getFollowAngle() - radar_frame->Azimuth) > this->radar_->getAngleDeviation()) {
                        if (radar_frame->Azimuth > 0)
                            // this->roclient_->SetVelocity(0, 0, -0.15);
                            z_vel = -0.15;
                        else if (radar_frame->Azimuth < 0) 
                            z_vel = 0.15;
                    }
                    else {
                        // this->roclient_->StopMove();
                        z_vel = 0;
                    }
                    if (std::abs(this->radar_->getFollowDistance() - radar_frame->Distance) > this->radar_->getDistanceDeviation()) {
                        if (radar_frame->Distance > this->radar_->getFollowDistance())
                            x_vel = 0.15;
                        else if (radar_frame->Distance < this->radar_->getFollowDistance())
                            x_vel = -0.15;
                    }
                    else {
                        // this->roclient_->StopMove();
                        x_vel = 0;
                    }
                    this->roclient_->SetVelocity(x_vel, 0, z_vel);
                }
                else {
                    this->roclient_->StopMove();
                }
            });

            radar_->Start();
            audio_.Start();
            
            thread_.clear();
            thread_.push_back(std::async(std::launch::async, &Application::POSIX_Message, this));
            thread_.push_back(std::async(std::launch::async, &Application::MQTT_Message, this));
            thread_.push_back(std::async(std::launch::async, &Application::Loop, this));
            // 断开连接
            mqttclient_->disconnect()->wait();
            std::cout << "Disconnected from MQTT broker." << std::endl;
            
        }
        catch (const std::exception& exc) {
            std::cerr << "Error: " << exc.what() << std::endl;
            return -1;
        }
    }
    return 0;       
}

bool Application::executeAndWaitForState(const std::string& target_state_str, std::function<void()> action_func, int timeout_ms) {
    if (action_func) {
        action_func(); // 执行动作
    }

    auto start_time = std::chrono::steady_clock::now();
    int sta;
    do {
        roclient_->GetFsmId(sta); 
        unitree::common::MilliSleep(100);

        // 检查超时
        auto current_time = std::chrono::steady_clock::now();
        if (std::chrono::duration_cast<std::chrono::milliseconds>(current_time - start_time).count() > timeout_ms) {
            std::cerr << "Timeout waiting for state: " << target_state_str << std::endl;
            return false; 
        }

        auto it = robot_fsm_map.find(sta);
        if (it != robot_fsm_map.end() && it->second == target_state_str) {
            return true;
        }
    } while (true);
    return false;
}

void Application::Parse_Json(std::string str) {
    auto robotctl = nlohmann::json::parse(str);
    auto CMD = robotctl["cmd"];
    // 处理 Unitree 机器人的数据
    if (CMD.contains("FORWARD")) {
        roclient_->SetVelocity(MOVE_VELOCITY, 0, 0, MOVE_DURATION);
    }

    if (CMD.contains("RETREAT")) {
        roclient_->SetVelocity(-MOVE_VELOCITY, 0, 0, MOVE_DURATION);
    }

    else if (CMD.contains("TRANSLATION_RIGHT")) {
        roclient_->SetVelocity(0, -MOVE_VELOCITY, 0, MOVE_DURATION);
    }

    else if (CMD.contains("TRANSLATION_LEFT")) {
        roclient_->SetVelocity(0, MOVE_VELOCITY, 0, MOVE_DURATION);
    }

    else if (CMD.contains("TWOHAND_KISS")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("two-hand kiss"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("KISS")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("left kiss"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("HANDS_UP")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("hands up"));
        });
    }

    else if (CMD.contains("CLAP")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("clap"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("HIGH_FIVE")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("high five"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("HUG")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("hug"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("REJECT")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("reject"));
            std::this_thread::sleep_for(std::chrono::seconds(5));
            armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("RIGHT_HAND_UP")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("right hand up"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("XRAY")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("x-ray"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("FACE_WAVE")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("face wave"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("HIGH_WAVE")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("high wave"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("SHAKE_HAND")) {
        Schedule([this]() {
            armclient_->ExecuteAction(armclient_->action_map.at("shake hand"));
            // std::this_thread::sleep_for(std::chrono::seconds(5));
            // armclient_->ExecuteAction(99);
        });
    }

    else if (CMD.contains("TURN_LEFT")) {
        Schedule([this]() {
            coordinates_.setTargetAngleZ(coordinates_.EularAngleZ() + (M_PI / 2.0f));
            roclient_->SwitchMoveMode(true);
            roclient_->Move(0, 0, MOVE_VELOCITY);
            {
                std::unique_lock<std::mutex> lock(coordinates_.mutex_);
                coordinates_.cv_.wait(lock);
            }
            roclient_->StopMove();
        });
    }

    else if (CMD.contains("TURN_RIGHT")) {
        Schedule([this]() {
            coordinates_.setTargetAngleZ(coordinates_.EularAngleZ() - (M_PI / 2.0f));
            roclient_->SwitchMoveMode(true);
            roclient_->Move(0, 0, -MOVE_VELOCITY);
            {
                std::unique_lock<std::mutex> lock(coordinates_.mutex_);
                coordinates_.cv_.wait(lock);
            }
            roclient_->StopMove();       
        });
    }

    else if (CMD.contains("TURN_LATER")) {
        Schedule([this]() {
            coordinates_.setTargetAngleZ(coordinates_.EularAngleZ() + M_PI);
            roclient_->SwitchMoveMode(true);
            roclient_->Move(0, 0, MOVE_VELOCITY);
            {
                std::unique_lock<std::mutex> lock(coordinates_.mutex_);
                coordinates_.cv_.wait(lock);
            }
            roclient_->StopMove();
        });
    }

    else if (CMD.contains("TURN")) {
        roclient_->Move(0, 0, MOVE_VELOCITY, true);
    }

    else if (CMD.contains("FOLLOW")) {
        if (CMD["FOLLOW"] == "true") {
            radar_->follow_mode_.store(true);
        }
        else if (CMD["FOLLOW"] == "false") {
            radar_->follow_mode_.store(false);
        }
    }

    else if(CMD.contains("damp")) {
        Schedule([this]() {
            roclient_->Damp();
            unitree::common::Sleep(10);
        });    
    }

    else if(CMD.contains("START")) {
        Schedule([this]() {
            int sta;
            roclient_->GetFsmId(sta);
            std::string current_fsm_sta;
            auto it = robot_fsm_map.find(sta);
            if (it != robot_fsm_map.end()) {
                current_fsm_sta = it->second;
            }
            else {
                std::cerr << "Invalid state: " << sta << std::endl;
                return;
            }
            if (current_fsm_sta == "ZeroTorque") {
                if (!executeAndWaitForState("Damp", [this]() {
                    roclient_->Damp();
                })) {
                    return;
                }
                current_fsm_sta = "Damp";
            }
            if (current_fsm_sta == "Damp" ){
                if (!executeAndWaitForState("StandUp", [this]() {
                    roclient_->StandUp();
                })) {
                    return;
                }
                std::cout << "Stand up!" << std::endl;
                current_fsm_sta = "StandUp";
            }
            if (current_fsm_sta == "StandUp") {
                std::this_thread::sleep_for(std::chrono::seconds(10));
                if (!executeAndWaitForState("Start", [this]() {
                    roclient_->Start();
                })) {
                    return;
                }
            }
            roclient_->GetFsmId(sta);
            it = robot_fsm_map.find(sta);
            if (it != robot_fsm_map.end() && it->second == "Start") {
                std::cout << "Robot is ready!" << std::endl;
                mqttclient_->publish("robot/status", "{\"status\":\"Start\"}", QOS, false);
            }
        });    
    }

    else if(CMD.contains("STOP")) {
        if (getFlag(FlagID::MASK_GAIT)) {
            roclient_->ContinuousGait(false);
            clearFlag(FlagID::MASK_GAIT);
        }
        else roclient_->StopMove();
        coordinates_.Relase();
        armclient_->ExecuteAction(99); // 停止手臂动作
    }
    
    else if(CMD.contains("HAND_SHAKE")) {
        Schedule([this]() {
            std::cout << "Shake hand Start...!" << std::endl;
            int32_t code = roclient_->ShakeHand(0);
            std::this_thread::sleep_for(std::chrono::seconds(5));
            roclient_->ShakeHand(1);
            std::cout << "Done!" << std::endl; 
            // code > 0? mqttclient_->publish("robot/status", "{\"shake_hand\":\"ok\"}", QOS, false): mqttclient_->publish("robot/status", "{\"shake_hand\":\"error\"}", QOS, false);
        });
    }

    else if(CMD.contains("HAND_WAVE")) {
        roclient_->WaveHand(false);
    }

    else if(CMD.contains("TURN_HAND_WAVE")) {
        roclient_->WaveHand(true);
    }

    else if(CMD.contains("squat")) {
    int32_t code = roclient_->Squat();
    code > 0? mqttclient_->publish("robot/status", "{\"squat\":\"ok\"}", QOS, false): mqttclient_->publish("robot/status", "{\"erro:ok\"}", QOS, false);
    
    }

    // else if(CMD.contains("stand_up")) {
    //     roclient_->StandUp();
    // }

    // else if(CMD.contains("sit")) {
    //     roclient_->Sit();
    // }

    else if(CMD.contains("stop_move")) {
        roclient_->StopMove();
    }

    else if(CMD.contains("MARK_TIME")) {
        if(CMD["MARK_TIME"] == "true") {
            roclient_->ContinuousGait(true);
            setFlag(FlagID::MASK_GAIT, true);
        }
    }
    else if(CMD.contains("Balance_Stand")) {
        roclient_->BalanceStand();
    }
    else if(CMD.contains("Get_status")) {
        int sta;
        roclient_->GetFsmId(sta);
        auto robot_status = nlohmann::json::object();
        if(robot_fsm_map.find(sta) == robot_fsm_map.end()) {
            std::cerr << "Invalid state: " << sta << std::endl;
            return;
        }
        auto it = robot_fsm_map.find(sta);
        if (it != robot_fsm_map.end()) {
            robot_status["status"] = it->second;
            std::cout << "Robot status: " << robot_fsm_map.at(sta) << std::endl;
            mqttclient_->publish("robot/status", robot_status.dump(), QOS, false);
        }
        else {
            std::cerr << "Invalid state: " << sta << std::endl;
            return;
        }
    }
    else if(CMD.contains("WAKE_STATE")) {
        if (CMD["WAKE_STATE"] == "AWAKE") {
            {
                std::unique_lock<std::mutex> lock(audio_.led_mutex_);
                audio_.wakeup_ = true;
                audio_.led_cv_.notify_one();
            }
        }
        else if (CMD["WAKE_STATE"] == "SLEPP") {
            {
                std::unique_lock<std::mutex> lock(audio_.led_mutex_);
                audio_.wakeup_ = false;
                audio_.led_cv_.notify_one();
            }
        }
    }

    // else if (CMD.contains("audio_merge")) {
    //     if(CMD["audio_merge"] == "Dialogue") {
    //         {
    //             std::unique_lock<std::mutex> lock(audio_.audio_mutex_);
    //             audio_.playsta.store(Play_Dialogue);
    //             audio_.audio_cv_.notify_one();
    //         }
    //     }
    //     else if(CMD["audio_merge"] == "Music") {
    //         {
    //             std::unique_lock<std::mutex> lock(audio_.audio_mutex_);
    //             audio_.playsta.store(Play_Music);
    //             audio_.audio_cv_.notify_one();
    //         }
    //     } 
    //     else if(CMD["audio_merge"] == "Stop") {
    //         {
    //             std::unique_lock<std::mutex> lock(audio_.audio_mutex_);
    //             audio_.playsta.store(Play_Stop);
    //             audio_.audio_cv_.notify_one();
    //         }
    //     }
    // }

    else if(CMD.contains("exit")) {
        roclient_->Damp();
        return;
    }
}